library(dplyr)
library(cregg)
library(tikzDevice)
library(ggplot2)

estimate_marginal_means <- function(experiment_data){
  # estimates marginal means for the second experiment
  formula_program <- num_prog ~ local_status + party
  formula_pork <- num_pork ~ local_status + party
  
  
  programmatic_outcome <- cregg::cj(
    formula = formula_program,
    data = experiment_data,
    id = ~respondent,
    estimate = "mm"
  )
  

  pork_outcome <- cregg::cj(
    formula = formula_pork,
    data = experiment_data,
    id = ~respondent,
    estimate = "mm"
  )
  
  second_experiment_result <- rbind(programmatic_outcome, pork_outcome)
  
  return(second_experiment_result)
  
}


format_cregg_output_for_plot <- function(second_experiment_result){
   # renames levels for the plot 
  
  second_experiment_result$outcome <- as.factor(second_experiment_result$outcome)
  levels(second_experiment_result$outcome) <- c(
    "Number of Particularistic Policies Selected", 
    "Number of Programmatic Policies Selected")
  
  levels(second_experiment_result$level) <- c("Elite", "Local", "Non-Local", "DPP", "LDP", "JCP", "CDPJ")
  
  return(second_experiment_result)
}

make_policy_type_selection_plot <- function(plot_data, main_paper_version = TRUE){
  
  # we do not show the party marginal means outside of the appendix, this 
  # filters it out of the plot 
  if (main_paper_version){
    plot_data <- plot_data %>% 
      filter(feature == "local_status")
  }
  
  
  policy_selection_plot <- ggplot(data = plot_data, aes(x = estimate, y = level, xmin = lower, xmax = upper)) +
    geom_point() + 
    geom_errorbar(width = 0) + 
    geom_vline(xintercept = 1.5, linetype = "dashed") + 
    facet_wrap(. ~ outcome, ncol = 1) + 
    scale_y_discrete(name = "Type of Candidate") + 
    scale_x_continuous(name = "Marginal Mean of Number of Policies Selected", limits = c(0.75, 2.)) +
    theme_bw()
  
  return(policy_selection_plot)
  
}


make_data_for_individual_items_comparison <- function(experiment_data){
  # this prepares data for plots that show the MMs by individual items in the
  # appendix 
  
  kigyoyuchi_formula <-  kigyo_yuchi ~ local_status + party
  kanko_sokushin_formula <-  kanko_sokushin ~ local_status + party
  kogyojigyo_formula <-  kogyojigyo ~ local_status + party
  seifusaimu_formula <-  seifusaimu ~ local_status + party
  gaiko_formula <-  gaiko ~ local_status + party
  nenkin_formula <-  nenkin ~ local_status + party
  
  call_cregg_for_given_formula <- function(x) cregg::cj(
    formula = x,
    data = experiment_data,
    id = ~respondent,
    estimate = "mm"
  )
  
  formulas <- list(kigyoyuchi_formula,
                   kanko_sokushin_formula,
                   kogyojigyo_formula,
                   seifusaimu_formula,
                   gaiko_formula,
                   nenkin_formula)
  
  cregg_objects <- lapply(formulas, call_cregg_for_given_formula)
  
  return(cregg_objects)
  
}

make_individual_policy_items_plot <- function(experiment_data){
  # makes the appendix plot for individual policy items 
  
  individual_item_mms <- make_data_for_individual_items_comparison(experiment_data)
  plot_data <- data.table::rbindlist(individual_item_mms)
  
  plot_data <- plot_data %>%
    filter(feature == "local_status")
  
  plot_data <- plot_data %>% 
    mutate(outcome = recode(
      outcome, 
      "gaiko" = "Foreign Affairs", 
      "kanko_sokushin" = "Promotion of Tourism",
      "kigyo_yuchi" = "Attraction of Businesses",
      "kogyojigyo" = "Public Works",
      "nenkin" = "National Pension Plan",
      "seifusaimu" = "National Debt"))
  
  pd <- position_dodge(0.1)
  
  ggplot(data = plot_data, aes(x = estimate, y = level, xmin = lower, xmax = upper, col = outcome)) +
    geom_point(position = pd) + 
    geom_errorbar(width = 0, position = pd) + 
    geom_vline(xintercept =0.5, linetype = "dashed") + 
    scale_y_discrete(name = "Type of Candidate") + 
    scale_x_continuous(name = "Marginal Mean of Policy Selection") +
    theme_bw()
}

path_to_data <- "data"
formatted_survey_file <- "follow_up_survey_formatted.rds"

path_to_figures <- "figures"

path_to_paper_figures <- "paper_plots"
mm_plot_file <- "num_policies_selected.eps"

path_to_appendix_figures <- "appendix_plots"

app_mm_plot_file <- "appendix_num_policies_selected.eps"
app_individual_items_plot_file <- "individual_policy_items_plot.eps"

formatted_data <- readRDS(file.path(path_to_data, formatted_survey_file))

estimated_mm <- estimate_marginal_means(formatted_data)
plot_data_formatted <- format_cregg_output_for_plot(estimated_mm)

policy_selection_plot <- make_policy_type_selection_plot(plot_data_formatted)
appendix_policy_selection_plot <- make_policy_type_selection_plot(plot_data_formatted,
                                                                  main_paper_version = FALSE)

individual_policy_items_plot <- make_individual_policy_items_plot(formatted_data)

ggsave(
  filename = file.path(path_to_figures, path_to_paper_figures, mm_plot_file),
  plot = policy_selection_plot,
  height = 8, 
  width = 6,
  units = "in",
  device = "eps")

ggsave(
  filename = file.path(path_to_figures, path_to_appendix_figures, app_mm_plot_file),
  plot = appendix_policy_selection_plot,
  height = 8, 
  width = 6,
  units = "in",
  device = "eps")

ggsave(
  filename = file.path(path_to_figures, path_to_appendix_figures, app_individual_items_plot_file),
  plot = individual_policy_items_plot,
  height = 8, 
  width = 6,
  units = "in",
  device = "eps")





